/* Copyright (c) 2021, National University of Defense Technology. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 * 
 *     http://www.apache.org/licenses/LICENSE-2.0
 * 
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "bsr_common.h"
#include "parcsr_matblockvec.h"

namespace YHAMG
{

#define MPI_TAG 150

void ParCSR_MatBlockVec(const ParCSR_Matrix& A, const Par_Vector& x, const Par_Vector& y)
{
	int n = A.LocLoc.size[0];
	int Ext_size = A.LocExt.size[1];

	int* A_LocLoc_rowptr = A.LocLoc.rowptr;
	int* A_LocLoc_colind = A.LocLoc.colind;
	double* A_LocLoc_values = A.LocLoc.values;

	int* A_LocExt_rowptr = A.LocExt.rowptr;
	int* A_LocExt_colind = A.LocExt.colind;
	double* A_LocExt_values = A.LocExt.values;

	if (A.LocLoc.size[0] == 0 && A.LocLoc.size[1] == 0) return;
	int block_size = A.LocLoc.size[0] == 0 ? x.Loc.size / A.LocLoc.size[1] : y.Loc.size / A.LocLoc.size[0];

	MPI_Comm comm = A.comm;
	int num_neighbors = A.num_neighbors;
	int* neighbor_ranks = A.neighbor_ranks;
	int* recv_ptr = A.Ext_ptr;
	int* send_ptr = A.send_ptr;
	int* send_list = A.send_list;
	double* send_buffer = new double[send_ptr[num_neighbors] * block_size];

	double* xv_Loc = x.Loc.values;
	double* yv_Loc = y.Loc.values;

	Vector x_Ext;
	x_Ext.Resize(Ext_size * block_size);
	double* xv_Ext = x_Ext.values;

	MPI_Request* recv_requests = new MPI_Request[num_neighbors];
	MPI_Request* send_requests = new MPI_Request[num_neighbors];
	MPI_Status status;

	for (int r = 0; r < num_neighbors; ++r)
		if (recv_ptr[r + 1] > recv_ptr[r])
			MPI_Irecv(xv_Ext + recv_ptr[r] * block_size, (recv_ptr[r + 1] - recv_ptr[r]) * block_size, MPI_DOUBLE, neighbor_ranks[r], MPI_TAG, comm, recv_requests + r);

	for (int r = 0; r < num_neighbors; ++r)
	{
		if (send_ptr[r + 1] > send_ptr[r])
		{
			for (int i = send_ptr[r]; i < send_ptr[r + 1]; ++i)
				VecBlockCopy(block_size, xv_Loc + send_list[i] * block_size, send_buffer + i * block_size);
			MPI_Isend(send_buffer + send_ptr[r] * block_size, (send_ptr[r + 1] - send_ptr[r]) * block_size, MPI_DOUBLE, neighbor_ranks[r], MPI_TAG, comm, send_requests + r);
		}
	}

#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
	for (int i = 0; i < n; ++i)
	{
		VecBlockFill(block_size, 0.0, yv_Loc + i * block_size);
		for (int j = A_LocLoc_rowptr[i]; j < A_LocLoc_rowptr[i + 1]; ++j)
			VecBlockScaleAdd(block_size, A_LocLoc_values[j], xv_Loc + A_LocLoc_colind[j] * block_size, yv_Loc + i * block_size);
	}

	for (int r = 0; r < num_neighbors; ++r)
		if (recv_ptr[r + 1] > recv_ptr[r])
			MPI_Wait(recv_requests + r, &status);

	for (int r = 0; r < num_neighbors; ++r)
		if (send_ptr[r + 1] > send_ptr[r])
			MPI_Wait(send_requests + r, &status);

#ifdef USE_OPENMP
#pragma omp parallel for schedule(guided)
#endif
	for (int i = 0; i < n; ++i)
		for (int j = A_LocExt_rowptr[i]; j < A_LocExt_rowptr[i + 1]; ++j)
			VecBlockScaleAdd(block_size, A_LocExt_values[j], xv_Ext + A_LocExt_colind[j] * block_size, yv_Loc + i * block_size);

	delete[] send_buffer;
	delete[] recv_requests;
	delete[] send_requests;
}

}